from tests.base_test import base_test_case
from app.utils.logger import get_logger

logger = get_logger(__name__)

TEST_DATA_DIR = base_test_case.test_data_dir
OUTPUT_DATA_DIR = base_test_case.output_data_dir
SUB_DIR = TEST_DATA_DIR / "1715339805571"

"""
使用paddleocr 封装的 PPStructureV3进行测试
测试发现 :PP-StructureV3 产线工具
    使用 from paddleocr import PPStructureV3与from paddlex import create_pipeline是一致的(参考测试脚本 PP-StructureV3_paddlex_test.py);
    说明:PaddleX 致力于实现产线级别的模型训练、推理与部署。模型产线是指一系列预定义好的、针对特定AI任务的开发流程，其中包含能够独立完成某类任务的单模型（单功能模块）组合
1 初始化性能测试
2 推理参数测试
3 可封装部分
"""


def paddleocr_with_ppstv3_pdf(input_file):
    from pathlib import Path
    from paddleocr import PPStructureV3

    output_path = Path(OUTPUT_DATA_DIR)

    pipeline = PPStructureV3()
    output = pipeline.predict(input=input_file)

    markdown_list = []
    markdown_images = []

    for res in output:
        md_info = res.markdown
        markdown_list.append(md_info)
        markdown_images.append(md_info.get("markdown_images", {}))

    markdown_texts = pipeline.concatenate_markdown_pages(markdown_list)

    mkd_file_path = output_path / f"{Path(input_file).stem}.md"
    mkd_file_path.parent.mkdir(parents=True, exist_ok=True)

    with open(mkd_file_path, "w", encoding="utf-8") as f:
        f.write(markdown_texts)

    for item in markdown_images:
        if item:
            for path, image in item.items():
                file_path = output_path / path
                file_path.parent.mkdir(parents=True, exist_ok=True)
                image.save(file_path)


if __name__ == '__main__':
    input_file = TEST_DATA_DIR / "1715339805571.pdf"
    # Not supported input data type! Only `numpy.ndarray` and `str` are supported! So has been ignored: /data/pyworks/fw_paddleocr_web/tests/test_files/1711605374231.pdf.
    paddleocr_with_ppstv3_pdf(str(input_file))
    pass
